#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@NAME          : sse.py
@TIME          : 2025/03/03 15:56:59
@AUTHOR        : chenlip
@VERSION       : 0.0.1
@DESCRIPTION   : 自定义SSE消息发布模块
'''
# --------------------------引用--------------------------
import queue
import threading
import json
import time
from flask import Response, stream_with_context, current_app
# --------------------------引用--------------------------

class MemorySSE:
    """基于内存的SSE实现，无需Redis"""
    
    # 通道映射: {channel_name: [{queue: Queue, last_active: timestamp}, ...]}
    _channels = {}
    _lock = threading.Lock()
    _last_cleanup_time = time.time()
    _cleanup_interval = 60  # 秒，避免过于频繁的清理
    _cleanup_running = False
    
    @classmethod
    def stream(cls, channel='global'):
        """创建SSE流"""
        message_queue = queue.Queue()

        # 添加时间戳，用于检测失效连接
        connection_data = {
            'queue': message_queue,
            'last_active': time.time()
        }
        
        # 注册队列到通道
        with cls._lock:
            if channel not in cls._channels:
                cls._channels[channel] = []
            cls._channels[channel].append(connection_data)  # 添加字典而非直接添加队列
        
        def event_stream():
            """生成SSE事件流"""
            yield 'event: connected\ndata: {"status":"connected"}\n\n'
            
            try:
                while True:
                    # 从队列获取消息
                    message = message_queue.get()
                    
                    # None表示结束信号
                    if message is None:
                        break

                    # 更新最后活动时间
                    connection_data['last_active'] = time.time()
                        
                    event_type = message.get('type', 'notification')
                    data = json.dumps(message)
                    
                    # 发送SSE格式的消息
                    yield f'event: {event_type}\ndata: {data}\n\n'
                    
            finally:
                # 连接关闭时清理
                with cls._lock:
                    if channel in cls._channels:
                        # 查找并移除当前连接
                        cls._channels[channel] = [
                            conn for conn in cls._channels[channel] 
                            if conn['queue'] != message_queue
                        ]
                        if not cls._channels[channel]:
                            del cls._channels[channel]
                            
        return Response(
            stream_with_context(event_stream()),
            mimetype='text/event-stream',
            headers={
                'Cache-Control': 'no-cache',
                'Connection': 'keep-alive',
                'X-Accel-Buffering': 'no'  # NGINX配置，禁用缓冲
            }
        )
    
    @classmethod
    def publish(cls, message, type='notification', channel='global'):
        """发布消息到指定通道"""
        with cls._lock:
            if channel not in cls._channels:
                current_app.logger.warning(f"尝试发布到不存在的通道: {channel}")
                return False
                
            # 构建消息
            msg = {
                'type': type,
                **message
            }
            
            # 发送到所有订阅者
            for conn in cls._channels[channel]:
                conn['queue'].put(msg)  # 这里正确访问字典中的queue
                
            current_app.logger.info(f"向通道 {channel} 发布了 {type} 类型的消息")
            # 检查是否需要进行清理
            cls._schedule_cleanup()
            
            return True
    
    @classmethod
    def get_user_channels(cls):
        """
        获取所有活跃的用户通道
        
        Returns:
            list: 用户通道列表
        """
        user_channels = []
        
        # 遍历所有通道，找出用户通道
        with cls._lock:
            for channel in cls._channels:  # 使用_channels而不是_clients
                # 假设用户通道的格式为 "user.{user_id}"
                if channel.startswith('user.') and len(cls._channels[channel]) > 0:
                    user_channels.append(channel)
        
        return user_channels
    
    @classmethod
    def _schedule_cleanup(cls):
        """安排清理任务"""
        current_time = time.time()
        
        # 检查是否需要执行清理
        if (current_time - cls._last_cleanup_time > cls._cleanup_interval and 
            not cls._cleanup_running):
            
            # 获取当前应用实例
            app = current_app._get_current_object()  # 重要：保存对应用对象的引用
            
            # 启动后台清理线程，传递应用实例
            cleanup_thread = threading.Thread(
                target=cls._cleanup_inactive_connections,
                args=(app,)  # 将应用实例传递给线程
            )
            cleanup_thread.daemon = True
            cleanup_thread.start()
            
            cls._last_cleanup_time = current_time
    
    @classmethod
    def _cleanup_inactive_connections(cls, app):
        """
        清理不活跃的连接
        
        Args:
            app: Flask应用实例
        """
        # 在线程内创建应用上下文
        with app.app_context():
            try:
                cls._cleanup_running = True
                current_app.logger.info("开始清理不活跃的SSE连接...")
                
                inactive_threshold = time.time() - 300  # 5分钟无活动视为失效
                cleaned_channels = 0
                cleaned_connections = 0
                
                with cls._lock:
                    channels_to_check = list(cls._channels.keys())
                    
                    for channel in channels_to_check:
                        if channel not in cls._channels:
                            continue
                            
                        # 过滤出活跃连接
                        active_connections = []
                        for conn in cls._channels[channel]:
                            if conn['last_active'] > inactive_threshold:
                                active_connections.append(conn)
                            else:
                                # 发送结束信号
                                try:
                                    conn['queue'].put(None)
                                except:
                                    pass
                                cleaned_connections += 1
                        
                        # 更新通道中的连接
                        if active_connections:
                            cls._channels[channel] = active_connections
                        else:
                            del cls._channels[channel]
                            cleaned_channels += 1
                
                current_app.logger.info(f"清理完成: 移除了 {cleaned_connections} 个不活跃连接, {cleaned_channels} 个空通道")
                
            except Exception as e:
                current_app.logger.error(f"清理SSE连接时出错: {str(e)}")
            finally:
                cls._cleanup_running = False